import taichi as ti
ti.init(arch=ti.x64) # only CPU backend is supported for now
 

n = 4
# step 1: create sparse matrix builder  构建稀疏矩阵
K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)  #max_num_triplets最多的非零项

@ti.kernel
def fill(A: ti.linalg.sparse_matrix_builder()):
    for i in range(n):
        A[i, i] += 1  # Only +=  and -= operators are supported for now.

# step 2: fill the builder with data.
fill(K)

print(">>>> K.print_triplets()")
K.print_triplets()
# outputs:
# >>>> K.print_triplets()
# n=4, m=4, num_triplets=4 (max=100)(0, 0) val=1.0(1, 1) val=1.0(2, 2) val=1.0(3, 3) val=1.0

# step 3: create a sparse matrix from the builder.
B = K.build()
print(">>>> B = K.build()")
print(B)
# outputs:
# >>>> B = K.build()
# [1, 0, 0, 0]
# [0, 1, 0, 0]
# [0, 0, 1, 0]
# [0, 0, 0, 1]